import os
import tensorflow as tf
from keras.datasets import mnist


def preprocess(x, y):
    """
    x is a simple image, not a batch
    """
    x = tf.cast(x, dtype=tf.float32) / 255.
    x = tf.reshape(x, [28 * 28])
    y = tf.cast(y, dtype=tf.int32)
    y = tf.one_hot(y, depth=10)
    return x, y


def load_mnist_data():
    """
        加载 MNIST 数据集数据
    """
    (train_images, train_labels), (test_images, test_labels) = mnist.load_data()

    print('datasets:', train_images.shape, test_images.shape)
    print('image min value:', train_images.min())
    print('image max value:', train_images.max())

    return (train_images, train_labels), (test_images, test_labels)


def process_mnist_data(x, y, x_val, y_val, batch_size=128):
    """
    处理 MNIST 数据集数据
    :param x:  训练图像
    :param y:   训练图像对应的标签
    :param x_val: 测试图像
    :param y_val: 测试图像对应的标签
    :param batch_size: 每批数据的大小
    :return:
    """
    db = tf.data.Dataset.from_tensor_slices((x, y))
    db = db.map(preprocess).shuffle(60000).batch(batch_size)
    ds_val = tf.data.Dataset.from_tensor_slices((x_val, y_val))
    ds_val = ds_val.map(preprocess).batch(batch_size)

    return db, ds_val


if __name__ == "__main__":
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

    (x, y), (x_val, y_val) = load_mnist_data()

    db, ds_val = process_mnist_data(x, y, x_val, y_val, 128)
    print(db)
    print(ds_val)

    sample = next(iter(db))
    print(sample[0].shape, sample[1].shape)